import json

from sqlalchemy import String, DateTime, JSON, TypeDecorator
from sqlalchemy.dialects.mysql import BIGINT, LONGBLOB
from sqlalchemy.orm import Mapped, mapped_column

from owl_admin.ext import db


class JSONEncoded(TypeDecorator):

    cache_ok = True

    impl = String

    def process_bind_param(self, value, dialect):
        if value is not None:
            return json.dumps(value, ensure_ascii=False)
        return None

    def process_result_value(self, value, dialect):
        if value is not None:
            return json.loads(value)
        return None


class AIChatFileMessagePO(db.Model):
    __tablename__ = 'ai_chat_file_message'
    __table_args__ = {'comment': 'AI聊天文件消息表'}

    id: Mapped[int] = mapped_column(BIGINT(20), primary_key=True, comment='主键')
    chat_id: Mapped[int] = mapped_column(BIGINT(20), comment='聊天ID')
    message_id: Mapped[int] = mapped_column(BIGINT(20), comment='消息ID')
    name: Mapped[str] = mapped_column(String(255), comment='文件名')
    path: Mapped[str] = mapped_column(String(255), comment='文件路径')
    suffix: Mapped[str] = mapped_column(String(255), comment='文件后缀')
    create_time: Mapped[DateTime] = mapped_column(DateTime, comment='创建时间')
    update_time: Mapped[DateTime] = mapped_column(DateTime, comment='更新时间')
    create_user: Mapped[int] = mapped_column(BIGINT(20), comment='创建人')
    update_user: Mapped[int] = mapped_column(BIGINT(20), comment='更新人')


class AIChatAgentPO(db.Model):
    __tablename__ = 'ai_chat_agent'
    __table_args__ = {'comment': 'AI聊天代理表'}

    id: Mapped[int] = mapped_column(BIGINT(20), primary_key=True, comment='主键')
    name: Mapped[str] = mapped_column(String(255), comment='代理名称')
    type: Mapped[str] = mapped_column(String(255), comment='代理类型')
    config: Mapped[str] = mapped_column(LONGBLOB, comment='代理配置')
    create_time: Mapped[DateTime] = mapped_column(DateTime, comment='创建时间')
    update_time: Mapped[DateTime] = mapped_column(DateTime, comment='更新时间')


class AIChatMessagePO(db.Model):
    __tablename__ = 'ai_chat_message'
    __table_args__ = {'comment': 'AI聊天消息表'}
    id: Mapped[int] = mapped_column(BIGINT(20), primary_key=True, comment='主键')
    chat_id: Mapped[int] = mapped_column(BIGINT(20), comment='聊天ID')
    role: Mapped[str] = mapped_column(String(255), comment='角色')
    content: Mapped[JSONEncoded] = mapped_column(JSONEncoded, comment='内容')
    files: Mapped[JSONEncoded] = mapped_column(JSONEncoded, comment='文件')
    create_time: Mapped[DateTime] = mapped_column(DateTime, comment='创建时间')
    update_time: Mapped[DateTime] = mapped_column(DateTime, comment='更新时间')
    create_user: Mapped[int] = mapped_column(BIGINT(20), comment='创建人')
    update_user: Mapped[int] = mapped_column(BIGINT(20), comment='更新人')


class AIChatPO(db.Model):
    __tablename__ = 'ai_chat'
    __table_args__ = {'comment': 'AI聊天表'}
    id: Mapped[int] = mapped_column(BIGINT(20), primary_key=True, comment='主键')
    name: Mapped[str] = mapped_column(String(255), comment='聊天名称')
    agent_id: Mapped[int] = mapped_column(BIGINT(20), comment='代理ID')
    create_time: Mapped[DateTime] = mapped_column(DateTime, comment='创建时间')
    update_time: Mapped[DateTime] = mapped_column(DateTime, comment='更新时间')
    create_by: Mapped[int] = mapped_column(BIGINT(20), comment='创建人')
    update_by: Mapped[int] = mapped_column(BIGINT(20), comment='更新人')
